import torch
from PIL import Image
from collections import Counter
from torchvision.transforms.functional import to_tensor, to_pil_image

"""
寻找背景颜色
@image Image对象（RGB模式）
"""
def find_bgc(image):

    w,h = image.size 
    color_count = {}

    for i in range(w):
        for j in range(h):
            pixel = image.getpixel((i,j))
            color_count[pixel] = color_count.get(pixel,0)+1

    most_common_color = Counter(color_count).most_common(1)

    try:
        return most_common_color[0][0]
    except Exception:
        return (0,0,0)

"""
替换背景图片背景颜色
@file 路径 或者 文件对象
@mode 使用CPU还是GPU
@image_type 生成图片类型 com->替换背景后的图片 pha抠出人像后的黑白轮廓图
@color 需要更换后的颜色 比如现在是红色要换成蓝色就是(0,0,255)
@now_color 现在的背景颜色（注意是现在的背景颜色不是需要换好后的背景颜色）
"""
def change_bgc(file,image_type="com",bg_color=(0,0,255),nowgb_color=None):

    src = Image.open(file)

    # 调整大小
    w = src.size[0] // 4 * 4
    h = src.size[1] //4 * 4 
    src = src.resize((w,h))
    # 调整模式
    if src.mode != "RGB":
        src = src.convert("RGB")
    # 创建背景图
    if not nowgb_color:
        nowgb_color = find_bgc(src)
    bgr = Image.new("RGB",(w,h),nowgb_color)
    
    mode = "cuda"
    # 查看是否有GPU
    try:
        src = to_tensor(src).cuda().unsqueeze(0)
        bgr = to_tensor(bgr).cuda().unsqueeze(0)
        model = torch.jit.load('./model_pth/bgc_change.pth').cuda().eval()
    except Exception as e:
        mode = "cpu"
        src = to_tensor(src).unsqueeze(0)
        bgr = to_tensor(bgr).unsqueeze(0)
        model = torch.jit.load('./model_pth/bgc_change.pth').eval()


    if src.size(2) <= 2048 and src.size(3) <= 2048:
        model.backbone_scale = 1/4
        model.refine_sample_pixels = 80_000
    else:
        model.backbone_scale = 1/8
        model.refine_sample_pixels = 320_000

    pha, fgr = model(src, bgr)[:2]
    com = pha * fgr + (1 - pha) * torch.tensor([bg_color[0]/255, bg_color[1]/255, bg_color[2]/255], device=mode).view(1, 3, 1, 1)

    if image_type == "pha":
        return to_pil_image(pha[0].cpu())
    else:
        return to_pil_image(com[0].cpu())


# change_bgc("src2.png","cpu").save("pha2.png")

